package com.leiyuee.djllearn;

import ai.djl.MalformedModelException;
import ai.djl.Model;
import ai.djl.basicdataset.cv.classification.Mnist;
import ai.djl.basicmodelzoo.basic.Mlp;
import ai.djl.inference.Predictor;
import ai.djl.metric.Metrics;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.transform.ToTensor;
import ai.djl.modality.cv.translator.ImageClassificationTranslator;
import ai.djl.ndarray.types.Shape;
import ai.djl.nn.Block;
import ai.djl.training.DefaultTrainingConfig;
import ai.djl.training.EasyTrain;
import ai.djl.training.Trainer;
import ai.djl.training.TrainingResult;
import ai.djl.training.dataset.Dataset;
import ai.djl.training.dataset.RandomAccessDataset;
import ai.djl.training.evaluator.Accuracy;
import ai.djl.training.listener.SaveModelTrainingListener;
import ai.djl.training.listener.TrainingListener;
import ai.djl.training.loss.Loss;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;
import ai.djl.translate.Translator;
import lombok.SneakyThrows;
import org.junit.jupiter.api.Nested;
import org.junit.jupiter.api.Test;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;


@Nested
public class DJLtest {
    //1.y要做人工智能 大模型 还是要python + 研究生

    //2.djl 是一个java的机器学习框架

    /**
     * 1. 准备数据集
     * 2. 构建神经网络
     * 3.创建模型(让这个模型应用前面的神经网络)
     * 4. 训练配置(如何训练)
     * 5. 保存模型
     * <p>
     * //
     * 6. 加载模型
     * 7. 预测(一个新输入-->判断这是什么)
     */

    @SneakyThrows
    private RandomAccessDataset getDataset(Dataset.Usage usage) throws Exception {
        //minst 是内置的,存在用户目录下的 ~/.djl/cache/
        Mnist mnist = Mnist.builder()
                .optUsage(usage)
                .setSampling(64, true)
                .build();
        mnist.prepare(new ProgressBar());
        return mnist;

        //自定义
//        ImageFolder build = ImageFolder.builder()
//                .setRepositoryPath("/path/to/your/dataset")
//                .setSampling(64, true)
//                .optImageHeight(28)
//                .optImageWidth(28)
//                .build();
    }

    @SneakyThrows
    @Test
    void trainingComplete() throws Exception {
        //1.准备数据集
        RandomAccessDataset trainingSet = getDataset(Dataset.Usage.TRAIN);
        RandomAccessDataset validateSet = getDataset(Dataset.Usage.TEST);

        //2.构建神经网络
        Block block = new Mlp(
                Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH,
                Mnist.NUM_CLASSES,
                new int[]{128, 64});
        //3.创建模型
        try (Model model = Model.newInstance("mlp")) {
            model.setBlock(block);
            //4.训练配置
            DefaultTrainingConfig config = setupTrainingConfig();
            //拿到训练器
            try (Trainer trainer = model.newTrainer(config)) {
                trainer.setMetrics(new Metrics());

                /*
                 * MNIST is 28x28 grayscale image and pre processed into 28 * 28 NDArray.
                 * 1st axis is batch axis, we can use 1 for initialization.
                 */
                Shape inputShape = new Shape(1, Mnist.IMAGE_HEIGHT * Mnist.IMAGE_WIDTH);

                // initialize trainer with proper input shape
                trainer.initialize(inputShape);

                //5、开始训练
                EasyTrain.fit(trainer, 5, trainingSet, validateSet);

                //6、训练结果
                TrainingResult result = trainer.getTrainingResult();
                System.out.println("result = " + result);
                //保存模型
                model.save(Paths.get("build/model"), "mlp"); //params  就是小模型文件了
            }
            }
        }



    private DefaultTrainingConfig setupTrainingConfig() throws Exception {
        String outputDir = "build/model";
        SaveModelTrainingListener listener = new SaveModelTrainingListener(outputDir);
        listener.setSaveModelCallback(
                trainer -> {
                    TrainingResult result = trainer.getTrainingResult();
                    Model model = trainer.getModel();
                    float accuracy = result.getValidateEvaluation("Accuracy");
                    model.setProperty("Accuracy", String.format("%.5f", accuracy));
                    model.setProperty("Loss", String.format("%.5f", result.getValidateLoss()));
                });

        return new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
                .addEvaluator(new Accuracy())
                .addTrainingListeners(TrainingListener.Defaults.logging(outputDir))
                .addTrainingListeners(listener);
    }


    //使用指定模型进行预测
    @Test
    void predictTest() throws IOException, MalformedModelException, TranslateException {
        //1.准备测试数据
        Image image = ImageFactory.getInstance().fromUrl("https://djl-ai.s3.amazonaws.com/resources/images/dog.jpg");
        //2.加载模型
        Path modelDir = Paths.get("build/model");
        Model model = Model.newInstance("mlp");
        model.setBlock(new Mlp(28 * 28, 10, new int[]{128, 64}));
        model.load(modelDir);
        //3.预测
        //如果没有 Translator，非标图像则无法处理
        //3.1 获取一个转换器
        Translator<Image, Classifications> translator =
                ImageClassificationTranslator.builder()
                        .addTransform(new Resize(Mnist.IMAGE_WIDTH, Mnist.IMAGE_HEIGHT))
                        .addTransform(new ToTensor())
                        .optApplySoftmax(true)
                        .build();
        //3.2 获取预测器
        Predictor<Image, Classifications> predictor = model.newPredictor(translator);
        //3.3 预测
        Classifications result = predictor.predict(image);
        System.out.println(result);
    }
}
